#pragma once

#include <cmath>
#include <vector>


using std::vector;

// 辅助函数：计算高斯权重
inline double gaussian(double x, double sigma)
{
    return exp(-(x * x) / (2 * sigma * sigma));
}

/**
 * @brief  双边滤波
 *
 * @tparam T
 * @param input 输入图像
 * @param diameter 滤波器直径（建议奇数）
 * @param sigmaColor 颜色空间标准差
 * @param sigmaSpace 空间域标准差
 * @return vector<vector<T>>
 */
template <typename T>
vector<vector<T>> bilateralFilter(
    const vector<vector<T>>& input,
    int diameter,
    double sigmaColor,
    double sigmaSpace
)
{
    if (input.empty() || input[0].empty()) {
        return {};
    }

    int half = diameter / 2;
    int rows = input.size();
    int cols = input[0].size();
    int channels = 1; // 默认灰度图像

    // 如果是彩色图像（如RGB），获取通道数
    if constexpr (std::is_same_v<T, vector<double>>) {
        channels = input[0].size();
    }

    // 预计算空间域权重（仅计算一次）
    vector<double> spaceWeights(diameter);
    double spaceSum = 0.0;
    for (int i = -half; i <= half; ++i) {
        double x = i;
        spaceWeights[half + i] = gaussian(x, sigmaSpace);
        spaceSum += spaceWeights[half + i];
    }

    // 归一化空间域权重
    for (auto& w : spaceWeights) {
        w /= spaceSum;
    }

    vector<vector<T>> output(rows, vector<T>(cols, 0));

    for (int i = 0; i < rows; ++i) {
        for (int j = 0; j < cols; ++j) {
            double sum = 0.0;
            double totalWeight = 0.0;

            // 遍历邻域
            for (int x = -half; x <= half; ++x) {
                for (int y = -half; y <= half; ++y) {
                    int xi = i + x;
                    int yj = j + y;

                    // 边界处理：超出边界的像素使用镜像填充
                    if (xi < 0) {
                        xi = -xi;
                    }
                    if (xi >= rows) {
                        xi = 2 * (rows - 1) - xi;
                    }
                    if (yj < 0) {
                        yj = -yj;
                    }
                    if (yj >= cols) {
                        yj = 2 * (cols - 1) - yj;
                    }

                    // 计算颜色域权重（灰度图像或彩色图像）
                    double colorDiff = 0.0;
                    if constexpr (std::is_same_v<T, vector<double>>) {
                        // 彩色图像：计算各通道差异的平方和
                        for (int c = 0; c < channels; ++c) {
                            double diff = (input[i][j][c] - input[xi][yj][c]);
                            colorDiff += diff * diff;
                        }
                    } else {
                        // 灰度图像：直接计算差值
                        colorDiff = pow(input[i][j] - input[xi][yj], 2);
                    }
                    double colorWeight = gaussian(colorDiff, sigmaColor);

                    // 总权重 = 空间权重 * 颜色权重
                    double weight = spaceWeights[x + half] * colorWeight;

                    // 累加加权和
                    sum += input[xi][yj] * weight;
                    totalWeight += weight;
                }
            }

            // 计算加权平均值
            output[i][j] = static_cast<T>(sum / totalWeight);
        }
    }

    return output;
}
